#include <mc_velocity_planning/dwa/dwa.h>

DWA::DWA(tf2_ros::Buffer* tf_buffer):VelocityPlanningBase(tf_buffer){
    // 参数配置  设置速度相关参数
    float __max_vel_vx = ros::param::param<float>("/dwa_planner/max_vel_vx", 0.8);
    float __min_vel_vx = ros::param::param<float>("/dwa_planner/min_vel_vx", 0.0);
    float __max_vel_vth = ros::param::param<float>("/dwa_planner/max_vel_vth", 1.2);
    float __min_vel_vth = ros::param::param<float>("/dwa_planner/min_vel_vth", -1.2);

    maxVel_[0] = min(__max_vel_limit[0], __max_vel_vx);
    maxVel_[2] = min(__max_vel_limit[2], __max_vel_vth);
    minVel_[0] = max(__min_vel_limit[0], __min_vel_vx);
    minVel_[2] = max(__min_vel_limit[2], __min_vel_vth);
    predictedTime_ = 3.0;
    cout << "[DWA] max_vel=" << msg_convert::Velocity2D2Str(maxVel_) << ", min_vel=" << msg_convert::Velocity2D2Str(minVel_)
        << ", predicted_time: " << predictedTime_ 
        << ", controller_frequency: " << __controller_frequency
        << endl;    

    // 配置ROS
    dwaPredictedTrajPub_ = nh_.advertise<sensor_msgs::PointCloud2>("/dwa_predicted_trajs", 1, true);
    dwaBestTrajPub_ = nh_.advertise<nav_msgs::Path>("/dwa_best_traj", 1, true);

    // 计算速度采样分辨率
    int vx_samples_nums = ceil((maxVel_[0] - minVel_[0]) / velSamples_[0]) + 1;
    int vth_samples_nums = ceil((maxVel_[2] - minVel_[2]) / velSamples_[2]) + 1;

    cout << "[DWA] vx_samples_nums: " << vx_samples_nums 
        << ", vth_samples_nums: " << vth_samples_nums 
        << endl;

    if(vx_samples_nums <=0 || vth_samples_nums <=0){
        ROS_ERROR("please check local_planner params!");
        vx_samples_nums = 1;
        vth_samples_nums = 1;
    }
    
    // 计算速度采样空间
    for(int i=0; i<vx_samples_nums; ++i){
        float vx = minVel_[0] + i * velSamples_[0];
        if(fabs(vx) < 0.01) continue;      //不计算0速
        if(vx > maxVel_[0]) break;
        for(int j=0; j<vth_samples_nums; ++j){
            float vth = minVel_[2] + j * velSamples_[2];
            if(vth > maxVel_[2]) break;
            velSpace_.emplace_back(Velocity2D(vx, 0.0, vth));

            Velocity2D vel_2d = Velocity2D(vx, 0.0, vth);
            cout << "[DWA] [ " << i << " | " << j << "], vel=" << msg_convert::Velocity2D2Str(vel_2d) << endl;
        }
    }

    // 计算预测轨迹 base_frame坐标系下
    for(auto vel : velSpace_){
        Trajectory traj;
        traj.vel_2d = vel;
        GenerateTrajectory(traj, predictedTime_, 0.1);
        predictedTraj_.emplace_back(traj);
    }
    PublishPointCloud2(predictedTraj_);
}


DWA::~DWA(){

}

// 计算轨迹和像素值
void DWA::GenerateTrajectory(Trajectory& traj, float predicted_time, float dt){
    // 计算预测点集数量
    int predicted_points = ceil(predicted_time / dt) + 1;
    Pose2D pose_2d;
    traj.waypoints.emplace_back(pose_2d);
    for(int i=0; i<predicted_points; ++i){
        // 检测时间
        if(i * dt > predicted_time) break;
        // 根据运动模型，计算轨迹  此处为差速运动模型
        pose_2d[0] += traj.vel_2d[0] * cos(pose_2d[2]) * dt;
        pose_2d[1] += traj.vel_2d[0] * sin(pose_2d[2]) * dt;
        pose_2d[2] += traj.vel_2d[2] * dt;
        // cout << "[GenerateTrajectory] i: " << i << ", pose_2d=" << pose_2d << endl;
        // 差距超过0.05时，记录一个像素
        // float dx = pose_2d.x - traj.waypoints.back().x;
        // float dy = pose_2d.y - traj.waypoints.back().y;
        // float dis = hypot(dx, dy);
        // if(dis > 0.04)
        traj.waypoints.emplace_back(pose_2d);
    }
}

// 发布点云数据
void DWA::PublishPointCloud2(vector<Trajectory> predicted_trajs){
    vector<geometry_msgs::Point> waypoints;
    for(auto traj : predicted_trajs){
        for(auto pose : traj.waypoints){
            geometry_msgs::Point point;
            point.x = pose[0];
            point.y = pose[1];
            point.z = traj.vel_2d[0];  //此处按速度vx显示
            waypoints.emplace_back(point);
        }
    }
    sensor_msgs::PointCloud2 cloud_msg;
    cloud_msg.header.frame_id = __base_frame;
    cloud_msg.header.stamp = ros::Time::now();
    cloud_msg.height = 1;
    cloud_msg.width = waypoints.size();
    cloud_msg.is_dense = false;

    sensor_msgs::PointCloud2Modifier modifier(cloud_msg);
    modifier.setPointCloud2FieldsByString(1, "xyz");
    modifier.resize(waypoints.size());

    // 使用迭代器填充点云消息的数据
    sensor_msgs::PointCloud2Iterator<float> iter_x(cloud_msg, "x");
    sensor_msgs::PointCloud2Iterator<float> iter_y(cloud_msg, "y");
    sensor_msgs::PointCloud2Iterator<float> iter_z(cloud_msg, "z");
    for(const auto& point : waypoints) {
        *iter_x = point.x;
        *iter_y = point.y;
        *iter_z = point.z;

        ++iter_x;
        ++iter_y;
        ++iter_z;
    }
    dwaPredictedTrajPub_.publish(cloud_msg);
}

// 发布最佳轨迹
void DWA::PublishBestTraj(Trajectory best_traj){
    nav_msgs::Path path_msg;
    path_msg.header.frame_id = __base_frame;
    path_msg.header.stamp = ros::Time::now();
    for(auto waypoint : best_traj.waypoints){
        geometry_msgs::PoseStamped pose_stamp;
        pose_stamp.pose = msg_convert::ToGeometryPose(waypoint);
        path_msg.poses.emplace_back(pose_stamp);
    }
    dwaBestTrajPub_.publish(path_msg);
}

// 计算刹车距离 v^2 - 0 = 2as
float DWA::StopDistance(Velocity2D vel_2d, float acc){
    return std::pow(vel_2d[0], 2) / (2 * acc);
}

// 计算航向角度评分  轨迹末端点角度和末端点与目标点连线的角度差
float DWA::HeadingScore(Pose2D traj_point, Pose2D goal){
    // 末端点与目标点连线的角度theta
    float goal_angle = math_utils::AnglePose2D(traj_point, goal);
    // 轨迹末端点角度 与  theta 的角度差
    float theta = math_utils::DifferenceAngle(traj_point[1], goal_angle);
    return 180.0 - fabs(math_utils::Rad2Deg(theta));
}

// 速度评分
float DWA::VelocityScore(Velocity2D vel_2d){
    return maxVel_[0] - vel_2d[0];     //速度越大越好
}

// 计算障碍物评分  计算轨迹末端点距离最近障碍物距离的评价函数 碰撞检测
bool DWA::ObstaclesScore(Pose2D traj_point, Velocity2D vel_2d, vector<Pose2D>& obstacles, 
    float inflation_radius, float& min_distance){
    // 计算轨迹末端点距离最近障碍物最小距离
    min_distance = std::numeric_limits<float>::infinity();
    for(auto& pose_2d : obstacles){
        float dis = math_utils::DistancePose2D(traj_point, pose_2d);
        dis -= inflation_radius;    //减去障碍物膨胀半径
        if(dis < 0){
            return false;
        }
        // cout << "[ObstaclesScore] dis: " << dis << endl;
        min_distance = min(min_distance, dis);
    }
    cout << "[ObstaclesScore] vel_2d: " << vel_2d << ", min_distance: " << min_distance << endl;

    //没有障碍物时，限制最大值， 此处设置为3 * inflation_radius
    min_distance = min(min_distance, 3 * inflation_radius);     
    return true;       
}

// 碰撞检测  简单的距离计算，后期优化
bool DWA::ObstaclesCollision(vector<Pose2D> waypoints, vector<Pose2D>& obstacles, 
    float inflation_radius){
    for(auto point : waypoints){
        for(auto obstacle : obstacles){
            float dis = math_utils::DistancePose2D(point, obstacle);
            if(dis < inflation_radius){
                return true;
            }
        }
    }
    return false;
}

// 查找最佳轨迹
bool DWA::FindBestTrajectory(Trajectory& best_traj, Pose2D local_goal){
    float inflation_radius = 0.2;   //膨胀半径
    float lidar_valid_range = 3.0;  //有效距离

    // 计算障碍物数据
    vector<Pose2D> obstacles;
    if(pointCloudXYZPtr_){
        for(auto point : pointCloudXYZPtr_->points){
            // 计算lidar_valid_range米内障碍物数据
            if(hypot(point.x, point.y) > lidar_valid_range) continue;
            obstacles.emplace_back(msg_convert::ToPose2D(point));
        }
        cout << "scan_data size: " << pointCloudXYZPtr_->points.size() 
            << ", obstacles size: " << obstacles.size() << endl; 
    }
    vector<Trajectory> valid_trajs;         //有效轨迹
    Trajectory sum_traj;    //用于归一化处理
    int best_i = -1;
    float best_traj_score = std::numeric_limits<float>::infinity();
    for(int i=0; i<predictedTraj_.size(); ++i){
        auto traj = predictedTraj_[i];
        if(traj.waypoints.empty()) continue;
        if(ObstaclesCollision(traj.waypoints, obstacles, inflation_radius)){
            continue;
        }
        Pose2D end_waypoint = traj.waypoints.back();
        // traj.obstacle_score = ObstaclesScore(end_waypoint, traj.vel_2d, obstacles, inflation_radius);
        if(!ObstaclesScore(end_waypoint, traj.vel_2d, obstacles, inflation_radius, traj.obstacle_score)) 
            continue;
        // 计算刹车距离
        float stop_distance = StopDistance(traj.vel_2d, 1.0);
        cout << "[ObstaclesScore] stop_distance: " << stop_distance << ",  traj.obstacle_score: " <<  traj.obstacle_score << endl;
        if(traj.obstacle_score < stop_distance){
            continue;
        }
        // // cout << "traj.obstacle_score: " << traj.obstacle_score << endl;
        // if(isinf(traj.obstacle_score)) continue;
        traj.goal_score = math_utils::DistancePose2D(end_waypoint, local_goal);
        traj.speed_score =  VelocityScore(traj.vel_2d);
        traj.heading_score = HeadingScore(end_waypoint, local_goal);
        traj.traj_score = traj.speed_score + traj.goal_score;
        if(traj.traj_score < best_traj_score){
            best_traj_score = traj.traj_score;
            best_i = i;
        }

        valid_trajs.emplace_back(traj);
        cout << "[debug | FindBestTrajectory] --> vel" << msg_convert::Velocity2D2Str(traj.vel_2d)
            << ", goal_score: " << traj.goal_score
            // << ", obstacle_score: " << traj.obstacle_score
            // << ", heading_score=" << traj.heading_score
            << ", speed_score=" << traj.speed_score
            << ", traj_score=" << traj.traj_score
            << endl;
        // 计算总分
        sum_traj.obstacle_score += traj.obstacle_score;
        sum_traj.speed_score += traj.speed_score;
        sum_traj.heading_score += traj.heading_score;
    }

    // for(int i=0; i<valid_trajs.size(); ++i){
    //     auto traj = valid_trajs[i];
    //     traj.obstacle_score =  traj.obstacle_score / std::max(sum_traj.obstacle_score, 0.01f);
    //     traj.heading_score = traj.heading_score / std::max(sum_traj.heading_score, 0.01f);
    //     traj.speed_score = traj.speed_score / std::max(sum_traj.speed_score, 0.01f);
    //     traj.traj_score = traj.obstacle_score + traj.heading_score + traj.speed_score;
    //     if(traj.traj_score < best_traj_score){
    //         best_traj_score = traj.traj_score;
    //         best_i = i;
    //     }
    // }   
    // 
    cout << "[debug | FindBestTrajectory] best_i=" << best_i << " "  << predictedTraj_.size() << endl;
    if(best_i < 0){
        cout << "[debug | FindBestTrajectory] can not find best trajectory!" << endl;
        return false;
    }
    best_traj = predictedTraj_[best_i];
    cout << "[debug | FindBestTrajectory] best_i=" << best_i
        << ", best_traj_score: " << best_traj_score
        << ", best_vel=" << msg_convert::Velocity2D2Str(best_traj.vel_2d)
        << endl;
    PublishBestTraj(best_traj);
    return true;
}

// 计算速度指令
bool DWA::ComputeTwist(Velocity2D curr_vel, Velocity2D& vel_2d, Pose2D local_goal){
    // 寻找最佳轨迹
    Trajectory best_traj;
    if(FindBestTrajectory(best_traj, local_goal)){
        vel_2d[0] = best_traj.vel_2d[0];
        // vel_2d.linear.y = best_traj.vel_2d.vy;
        vel_2d[2] = best_traj.vel_2d[2];
        return true;
    }
    return false;
}